Greptile Summary

This PR adds a _dtype field (the "fake" high-precision dtype) to every QuantizedTensorStorage subclass and threads it through constructors, get_metadata(), and dequantize() defaults. This cleanly removes the hardcoded torch.bfloat16 guesses in distributed.py and ensures that dequantize() called without arguments no longer silently degrades to FP32.

Key changes and concerns:

• fake_dtype validation gate in QuantizedTensor.__new__ — The guard if fake_dtype is not None and fake_dtype != dtype: raise ValueError(...) (line 376–377 of quantized_tensor.py) is logically correct in isolation, but since get_metadata() now always emits fake_dtype=self._dtype, any code that calls make_like(tensor, dtype=new_dtype) with a different dtype will immediately hit this ValueError. There are 10+ existing call sites in the codebase (attention/dot_product_attention/context_parallel.py, utils.py, tensor/__init__.py) that pass an explicit dtype to make_like, and the module-level cast_to_dtype helper (used by model.half() / model.bfloat16()) iterates over all QuantizedTensor parameters using exactly this pattern. This is a regression that will surface in normal training workflows.
• No backward-compatibility guard for old pickled storage objects — The new _dtype field is only set in __new__; there is no hasattr fallback in dequantize() (unlike the dtype property on QuantizedTensor). Unpickling a *TensorStorage that was saved before this PR will result in AttributeError: _dtype on the first dequantize call.
• Storage classes correctly updated — All four storage classes (Float8TensorStorage, Float8BlockwiseQTensorStorage, MXFP8TensorStorage, NVFP4TensorStorage) follow a consistent pattern and correctly dispatch between storage-only (object.__new__) and full-tensor (super().__new__) paths.

Confidence Score: 2/5

• Not safe to merge as-is — the fake_dtype validation in QuantizedTensor.__new__ will break existing make_like(dtype=X) call sites throughout the attention module and the module-level dtype-cast utility.
• The core idea (storing the high-precision dtype on storage objects) is sound and the distributed.py and C++ changes are clean improvements. However, the new validation guard in QuantizedTensor.__new__ combined with fake_dtype being unconditionally included in get_metadata() creates a regression: any call to make_like or to_dtype that changes the nominal dtype — including the widely-used cast_to_dtype module helper — will now raise a ValueError. This is a blocking correctness issue for normal training workflows.
• transformer_engine/pytorch/quantized_tensor.py — the fake_dtype validation logic and its interaction with get_metadata() requires the most attention before merging.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["C++ Quantizer::create_tensor(dtype)"] -->|"internal=true"| B["Float8TensorStorage.__new__\n(fake_dtype=GetATenDType(dtype))"]
    A -->|"internal=false"| C["Float8Tensor.__new__\n(dtype=GetATenDType(dtype))"]

    B -->|"cls is Storage"| D["object.__new__()\ninstance._dtype = fake_dtype"]
    B -->|"cls is Tensor subclass"| E["super().__new__(cls, fake_dtype=fake_dtype)"]
    C --> E

    E --> F["QuantizedTensor.__new__(dtype, fake_dtype)\nValidate: fake_dtype == dtype if not None\ninstance._dtype = dtype"]

    F --> G["QuantizedTensor._dtype set"]
    D --> G

    G -->|"dequantize() called"| H{"dtype arg?"}
    H -->|"None"| I["use self._dtype"]
    H -->|"explicit"| J["use explicit dtype"]
    I --> K["dequantize to correct high-precision dtype"]
    J --> K

    style D fill:#90EE90
    style G fill:#90EE90
    style F fill:#FFB6C1,stroke:#FF0000
    note1["⚠️ Validation fails for\nmake_like(dtype=X) when X != _dtype"]
    F -.-> note1

_{Last reviewed commit: 369f8b5}

/te-ci pytorch

/te-ci pytorch

/te-ci pytorch